import torch
from torchvision import transforms
import numpy as np
from PIL import Image
from numpy.lib.stride_tricks import sliding_window_view
from DCP import rgb_to_dark_channel


# ---------------------- 函数1：估计大气光值 ----------------------
def estimate_atmospheric_light(image_rgb, dark_channel,window_size,top_percent=0.1):#py中默认形参后不可跟随非默认形参
    """
    根据暗通道图估计大气光值A
    :param image_rgb: 原始RGB图像，形状[H, W, 3]，值域0-255,原始RGB图像的numpy数组
    :param dark_channel: 暗通道图，形状[H, W]，暗通道图像的numpy数组
    :param top_percent: 选择最亮像素的百分比（单位：%，默认0.1%）
    :return: 大气光值A，形状[3,]的numpy数组 这个形状表示的是一个一维数组，内含元素为3，numpy的一维数组尺寸就是这样表示的，别问
    """
    """
    这里算大气光值的思路是将暗通道的numpy数组展平成一维numpy数组
    然后通过排序找到索引
    再切换回原始图像中找每个区域最亮的点然后求均值
    """
    # 将暗通道展平
    h, w = dark_channel.shape #这个函数返回的是一个元组，元组当然其内不一定是只能有两个元素哦
    flat_dark = dark_channel.ravel() #这个方法是将一个多维数组展平为一个一维数组
    num_pixels = flat_dark.size #size返回的是数组中元素个数，shape返回的是数组的结构
  #这里为什么要展平呢？
    # 计算需要保留的像素数
    num_top = int(num_pixels * top_percent / 100) #这里是算要取暗通道图中几个像素点，原始的是0.1%，所以要除100
    if num_top < 1: #如果算出来像素点数小于1，那至少将像素数量设置为1。
        num_top = 1  # 至少保留一个像素
    # 获取最亮像素的索引
    indices = np.argpartition(flat_dark, -num_top)[-num_top:]
    #上面这个np.argpartition有点意思，我们仔细看一下，这个函数会将数组进行自动排序左边是k个小的值，右边是剩余的元素，但是值都比左边大，下面是一个示例
    #这个函数的形参第一个是数组，第二个是排序的元素个数，也就是将多少个最小的元素找出来放在左边
    #但是在这里这个排序元素个数是负的，这里是就几个最大的元素找出来放在右边
    #会注意后面有一个切片，因为那个函数返回的本来就是一个数组，这个切片是把符合条件的几个最大值拿出来
    """
    arr = np.array([3, 1, 4, 2, 5])
    indices = np.argpartition(arr, 2)  # 找到最小的2个元素的索引
    print(indices)  # 输出: [1 3 0 2 5]（前两个索引对应最小的两个元素1和2）输出会发现是一个原索引组成的数组
    """
    # 将索引转换为原图坐标
    rows, cols = np.unravel_index(indices, (h, w))
    #woc这个函数牛逼，第一个参数是展平后一维数组中的索引，第二个参数是原数组形状，最后返回的结果是一个元组(rows, cols)，rows中包含所有找的点的行索引，另一个包含所有列索引
    # 初始化累积最大值
    A_sum = np.zeros(3, dtype=np.float64) #创建一个含有三个元素的一维数组，元素类型为float64
    # 计算窗口半径
    radius = window_size // 2 #这里估计又要开始滑动读取了。
    # 取每个最亮像素所在区域计算最大值并累积
    for row, col in zip(rows, cols): #这个zip有点牛逼，会生成一个个对应的元组
        # 计算区域边界（防止越界）
        row_min = max(0, row - radius)
        row_max = min(h, row + radius + 1)
        col_min = max(0, col - radius)
        col_max = min(w, col + radius + 1)
        #这上面这个计算区域边界感觉我一辈子都写不出来
        # 提取区域并计算最大值 这边我已经有点看不懂了，但是就是找点然后取最大值
        region = image_rgb[row_min:row_max, col_min:col_max, :]
        region_max = np.max(region, axis=(0, 1))  # 每个通道的最大值
        # 累积
        A_sum += region_max
    # 计算平均值
    A = A_sum / num_top
    return A.astype(np.uint8)
# ---------------------- 函数2：估计透射率（含ω参数）---------------------
def estimate_transmission(dark_channel, A, omega=0.95, t_min=0.1, t_max=0.9):#omega这个是透射率前面的未知参数哦
    """
    估计透射率图
    :param dark_channel: 暗通道图，形状[H, W]
    :param A: 大气光值，形状[3,]的numpy数组
    :param omega: 透射率调节参数（默认0.95）
    :param t_min: 透射率下限（默认0.1）
    :param t_max: 透射率上限（默认0.9） 这边搞两个透射率上下限说实话我没看懂在干嘛
    :return: 透射率图，形状[H, W]
    """
    # 将暗通道归一化到[0,1]
    dark_norm = dark_channel.astype(np.float32) / 255.0 #这里是防止算着算着数值溢出了
    # 计算A的灰度值（取各通道均值）
    A_gray = np.mean(A) / 255.0 #它这个实现是简化后的实现，正常是要求各通道和对应大气光值对应最小值，这里直接把大气光值从rgb变为灰度图
    # 计算原始透射率
    t = 1 - omega * (dark_norm / A_gray)
    # 约束透射率范围
    t = np.clip(t, t_min, t_max)
    return t
#--------------------------------更严格的透射率估计的实现----------------------------
def compute_transmission_strict(image_rgb, A, window_size=15, omega=0.95):
    """
    严格按照公式计算透射率:
    t^wk = 1 - ω * [min_{c∈{r,g,b}} min_{x∈wk} (I^c(x)/A^c)]

    :param image_rgb: 输入RGB图像, 形状[H, W, 3], 值域0-255
    :param A: 大气光值, 形状[3,]的数组
    :param window_size: 窗口大小(奇数)
    :param omega: 透射率调节参数
    :return: 透射率图, 形状[H, W]
    """
    # 参数校验
    window_size = window_size if window_size % 2 else window_size + 1
    radius = window_size // 2
    h, w, _ = image_rgb.shape
    # 转换为float类型避免溢出
    I = image_rgb.astype(np.float32)
    A = A.astype(np.float32) + 1e-10  # 避免除零
    # Step 1: 计算每个像素的I^c(x)/A^c
    ratio_map = I / A[None, None, :]
    # Step 2: 对每个颜色通道进行窗口最小值计算
    padded_ratio = np.pad(ratio_map,
                          ((radius, radius), (radius, radius), (0, 0)),
                          mode='reflect')
    # 创建滑动窗口视图 (H, W, C, window_size, window_size)
    windows = sliding_window_view(padded_ratio, (window_size, window_size), axis=(0, 1))
    # 计算每个窗口的最小值 (H, W, C)
    min_per_channel = np.min(windows, axis=(3, 4))
    # Step 3: 跨颜色通道取最小值 (H, W)
    min_over_channels = np.min(min_per_channel, axis=2)
    # Step 4: 应用ω参数计算透射率
    t = 1 - omega * min_over_channels
    # 约束透射率范围
    t = np.clip(t, 0.1, 0.99) #约束t的值，不在这个值中的t会被自动设置为最接近的边界。
    return t
# ---------------------- 函数3：大气散射模型逆运算 ----------------------
def apply_scattering_model(hazy_image, t, A, t0=0.1):
    """
    根据大气散射模型恢复无雾图像
    :param hazy_image: 有雾图像，形状[H, W, 3]
    :param t: 透射率图，形状[H, W]
    :param A: 大气光值，形状[3,]
    :param t0: 透射率最小值保护阈值（默认0.1）
    :return: 去雾后的图像，形状[H, W, 3]
    """
    # 转换数据类型
    hazy = hazy_image.astype(np.float32)
    A = A.astype(np.float32)
    # 扩展t到3通道
    t_expanded = np.expand_dims(t, axis=2)
    t_expanded = np.maximum(t_expanded, t0)  # 避免除以零
    # 计算去雾图像
    J = (hazy - A) / t_expanded + A
    # 约束值域并转换类型
    J = np.clip(J, 0, 255).astype(np.uint8)
    return J
#上面的我懒得仔细看了，又不是我看了一遍就能学会的。接下来直接开始测试
if __name__ == '__main__':
    image_path='D:\dehaze\\real_haze3.jpg'
    patch_size=15
    omega=0.95
    image = Image.open(image_path).convert('RGB')
    transform=transforms.Compose([
        transforms.Resize((512,512))
    ])
    image=transform(image)
    image=np.array(image)
    dark_channel=rgb_to_dark_channel(image,patch_size)
    #估计大气光值的参数是原图像，暗通道图像，局部区域大小和取点百分比
    A=estimate_atmospheric_light(image,dark_channel,patch_size)
    #透射率严格估计的参数是原图像，大气光值，窗口大小，透射率调节参数
    t_real = compute_transmission_strict(image, A, patch_size, omega)
    #透射率的简单估计的参数为暗通道数组和大气光值
    t_simple=estimate_transmission(dark_channel,A,omega)
    haze_image=Image.fromarray(image)
    haze_image.show()
    dehaze_image1=apply_scattering_model(image,t_simple,A)
    dehaze_image2 = apply_scattering_model(image, t_real, A)
    dehaze_image1=Image.fromarray(dehaze_image1)
    dehaze_image2=Image.fromarray(dehaze_image2)
    dehaze_image1.show()
    dehaze_image2.show()
#说实话简化的透射率求法和严格的求法得到的无雾图像差别我肉眼真没看出来多大
#但是去完雾的图片有一圈白色边缘倒是奇怪